{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于采样的路径规划算法——RRM\n",
    "\n",
    "\n",
    "博客参考：https://blog.csdn.net/weixin_42301220/article/details/125254296\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.spatial import KDTree\n",
    "from celluloid import Camera  # 保存动图时用，pip install celluloid\n",
    "%matplotlib qt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter\n",
    "N_SAMPLE = 500  # 采样点数目，即随机点集V的大小\n",
    "N_KNN = 10  # 一个采样点的领域点个数\n",
    "MAX_EDGE_LEN = 30.0  # [m] Maximum edge length\n",
    "\n",
    "show_animation = True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "声明节点，使用图搜索算法（这里是迪杰斯特拉算法）时使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "    \"\"\"\n",
    "    Node class for dijkstra search\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, x, y, cost, parent_index):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "        self.cost = cost  # 每条边权值\n",
    "        self.parent_index = parent_index\n",
    "\n",
    "    def __str__(self):\n",
    "        return str(self.x) + \",\" + str(self.y) + \",\" +\\\n",
    "            str(self.cost) + \",\" + str(self.parent_index)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### prm规划入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prm_planning(start_x, start_y, goal_x, goal_y, obstacle_x_list, obstacle_y_list, robot_radius, *, camara=None, rng=None):\n",
    "    \"\"\"\n",
    "    Run probabilistic road map planning\n",
    "\n",
    "    :param start_x: start x position\n",
    "    :param start_y: start y position\n",
    "    :param goal_x: goal x position\n",
    "    :param goal_y: goal y position\n",
    "    :param obstacle_x_list: obstacle x positions\n",
    "    :param obstacle_y_list: obstacle y positions\n",
    "    :param robot_radius: robot radius\n",
    "    :param rng: 随机数构造器\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    obstacle_kd_tree = KDTree(np.vstack((obstacle_x_list, obstacle_y_list)).T)\n",
    "    # 采样点集生成\n",
    "    sample_x, sample_y = sample_points(start_x, start_y, goal_x, goal_y,\n",
    "                                       robot_radius,\n",
    "                                       obstacle_x_list, obstacle_y_list,\n",
    "                                       obstacle_kd_tree, rng)\n",
    "    if show_animation:\n",
    "        plt.plot(sample_x, sample_y, \".b\")\n",
    "\n",
    "    # 生成概率路图\n",
    "    road_map = generate_road_map(\n",
    "        sample_x, sample_y, robot_radius, obstacle_kd_tree)\n",
    "    # 使用迪杰斯特拉规划路径\n",
    "    rx, ry = dijkstra_planning(\n",
    "        start_x, start_y, goal_x, goal_y, road_map, sample_x, sample_y, camara)\n",
    "\n",
    "    return rx, ry\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 判断碰撞"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def is_collision(sx, sy, gx, gy, rr, obstacle_kd_tree):\n",
    "    \"\"\"判断是否发生碰撞,true碰撞，false不碰\n",
    "        rr: 机器人半径\n",
    "    \"\"\"\n",
    "    x = sx\n",
    "    y = sy\n",
    "    dx = gx - sx\n",
    "    dy = gy - sy\n",
    "    yaw = math.atan2(gy - sy, gx - sx)\n",
    "    d = math.hypot(dx, dy)\n",
    "\n",
    "    if d >= MAX_EDGE_LEN:\n",
    "        return True\n",
    "\n",
    "    D = rr\n",
    "    n_step = round(d / D)\n",
    "\n",
    "    for i in range(n_step):\n",
    "        dist, _ = obstacle_kd_tree.query([x, y])  # 查询kd-tree附近的邻居\n",
    "        if dist <= rr:\n",
    "            return True  # collision\n",
    "        x += D * math.cos(yaw)\n",
    "        y += D * math.sin(yaw)\n",
    "\n",
    "    # goal point check\n",
    "    dist, _ = obstacle_kd_tree.query([gx, gy])\n",
    "    if dist <= rr:\n",
    "        return True  # collision\n",
    "\n",
    "    return False  # OK\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 概率路图的构建过程的伪代码如下：\n",
    "\n",
    "<center><img src=\"https://img-blog.csdnimg.cn/c67b48b9170747b49f1d4f48238543aa.png\" width=40%>\n",
    "\n",
    "\n",
    "如上图所示，\n",
    "\n",
    "- 步骤 1、2 : 初始化两个集合，其中$V$表示随机点集，$E$表示路径集。\n",
    "\n",
    "\n",
    "- 步骤 3~8 : 每次随机采样一个无碰撞的点，将这个无碰撞的点加入$V$中，重复$n$次。\n",
    "\n",
    "- 步骤 9~16：生成概率路图。\n",
    "\t- 步骤 10：对$V$中的每个点$q$，根据一定的距离范围选择$k$个邻域点\n",
    "\t- 步骤 11~15：对每个领域点$q'$进行判断，如果$q$和$q'$尚未形成路径，则将其连接形成路径，随后进行碰撞检测，若无碰撞，则保留该路径。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_road_map(sample_x, sample_y, rr, obstacle_kd_tree):\n",
    "    \"\"\"\n",
    "    概率路图生成\n",
    "\n",
    "    sample_x: [m] x positions of sampled points\n",
    "    sample_y: [m] y positions of sampled points\n",
    "    robot_radius: Robot Radius[m]\n",
    "    obstacle_kd_tree: KDTree object of obstacles\n",
    "    \"\"\"\n",
    "\n",
    "    road_map = []\n",
    "    n_sample = len(sample_x)\n",
    "    sample_kd_tree = KDTree(np.vstack((sample_x, sample_y)).T)\n",
    "\n",
    "    for (i, ix, iy) in zip(range(n_sample), sample_x, sample_y):\n",
    "        # 对V中的每个点q，选择k个邻域点\n",
    "        dists, indexes = sample_kd_tree.query([ix, iy], k=n_sample)\n",
    "        edge_id = []\n",
    "\n",
    "        for ii in range(1, len(indexes)):\n",
    "            nx = sample_x[indexes[ii]]\n",
    "            ny = sample_y[indexes[ii]]\n",
    "            # 对每个领域点$q'$进行判断，如果$q$和$q'$尚未形成路径，则将其连接形成路径并进行碰撞检测，若无碰撞，则保留该路径。\n",
    "            if not is_collision(ix, iy, nx, ny, rr, obstacle_kd_tree):\n",
    "                edge_id.append(indexes[ii])\n",
    "\n",
    "            if len(edge_id) >= N_KNN:\n",
    "                break\n",
    "\n",
    "        road_map.append(edge_id)\n",
    "\n",
    "    #  plot_road_map(road_map, sample_x, sample_y)\n",
    "\n",
    "    return road_map\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 迪杰斯特拉算法实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dijkstra_planning(sx, sy, gx, gy, road_map, sample_x, sample_y, camara):\n",
    "    \"\"\"\n",
    "    s_x: start x position [m]\n",
    "    s_y: start y position [m]\n",
    "    goal_x: goal x position [m]\n",
    "    goal_y: goal y position [m]\n",
    "    obstacle_x_list: x position list of Obstacles [m]\n",
    "    obstacle_y_list: y position list of Obstacles [m]\n",
    "    robot_radius: robot radius [m]\n",
    "    road_map: 构建好的路图 [m]\n",
    "    sample_x: 采样点集x [m]\n",
    "    sample_y: 采样点集y [m]\n",
    "\n",
    "    @return: Two lists of path coordinates ([x1, x2, ...], [y1, y2, ...]), empty list when no path was found\n",
    "    \"\"\"\n",
    "\n",
    "    start_node = Node(sx, sy, 0.0, -1)\n",
    "    goal_node = Node(gx, gy, 0.0, -1)\n",
    "    # 使用字典的方式构造开闭集合\n",
    "    # openList表由待考察的节点组成， closeList表由已经考察过的节点组成。\n",
    "    open_set, closed_set = dict(), dict()\n",
    "    open_set[len(road_map) - 2] = start_node\n",
    "\n",
    "    path_found = True\n",
    "    # 步骤与A星算法一致\n",
    "    while True:\n",
    "        # 如果open_set是空的\n",
    "        if not open_set:\n",
    "            print(\"Cannot find path\")\n",
    "            path_found = False\n",
    "            break\n",
    "\n",
    "        c_id = min(open_set, key=lambda o: open_set[o].cost)\n",
    "        current = open_set[c_id]\n",
    "\n",
    "        # show graph\n",
    "        if show_animation and len(closed_set.keys()) % 2 == 0:\n",
    "            # for stopping simulation with the esc key.\n",
    "            plt.gcf().canvas.mpl_connect(\n",
    "                'key_release_event',\n",
    "                lambda event: [exit(0) if event.key == 'escape' else None])\n",
    "            plt.plot(current.x, current.y, \"xg\")\n",
    "            plt.pause(0.001)\n",
    "            if camara != None:\n",
    "                camara.snap()\n",
    "\n",
    "        if c_id == (len(road_map) - 1):\n",
    "            print(\"goal is found!\")\n",
    "            goal_node.parent_index = current.parent_index\n",
    "            goal_node.cost = current.cost\n",
    "            break\n",
    "\n",
    "        # Remove the item from the open set\n",
    "        del open_set[c_id]\n",
    "        # Add it to the closed set\n",
    "        closed_set[c_id] = current\n",
    "\n",
    "        # expand search grid based on motion model\n",
    "        for i in range(len(road_map[c_id])):\n",
    "            n_id = road_map[c_id][i]\n",
    "            dx = sample_x[n_id] - current.x\n",
    "            dy = sample_y[n_id] - current.y\n",
    "            d = math.hypot(dx, dy)\n",
    "            node = Node(sample_x[n_id], sample_y[n_id],\n",
    "                        current.cost + d, c_id)\n",
    "\n",
    "            if n_id in closed_set:\n",
    "                continue\n",
    "            # Otherwise if it is already in the open set\n",
    "            if n_id in open_set:\n",
    "                if open_set[n_id].cost > node.cost:\n",
    "                    open_set[n_id].cost = node.cost\n",
    "                    open_set[n_id].parent_index = c_id\n",
    "            else:\n",
    "                open_set[n_id] = node\n",
    "\n",
    "    if path_found is False:\n",
    "        return [], []\n",
    "\n",
    "    # generate final course\n",
    "    rx, ry = [goal_node.x], [goal_node.y]\n",
    "    parent_index = goal_node.parent_index\n",
    "    while parent_index != -1:\n",
    "        n = closed_set[parent_index]\n",
    "        rx.append(n.x)\n",
    "        ry.append(n.y)\n",
    "        parent_index = n.parent_index\n",
    "\n",
    "    return rx, ry\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 采样点集生成\n",
    "\n",
    "即集合V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_points(sx, sy, gx, gy, rr, ox, oy, obstacle_kd_tree, rng):\n",
    "    \"\"\"采样点集生成\n",
    "    \"\"\"\n",
    "    max_x = max(ox)\n",
    "    max_y = max(oy)\n",
    "    min_x = min(ox)\n",
    "    min_y = min(oy)\n",
    "\n",
    "    sample_x, sample_y = [], []\n",
    "\n",
    "    if rng is None:\n",
    "        rng = np.random.default_rng()\n",
    "\n",
    "    while len(sample_x) <= N_SAMPLE:\n",
    "        tx = (rng.random() * (max_x - min_x)) + min_x\n",
    "        ty = (rng.random() * (max_y - min_y)) + min_y\n",
    "\n",
    "        # 在障碍物中查询离[tx, ty]最近的点的距离\n",
    "        dist, index = obstacle_kd_tree.query([tx, ty])\n",
    "\n",
    "        # 距离大于机器人半径，说明没有碰撞，将这个无碰撞的点加入V中，重复n次。\n",
    "        if dist >= rr:\n",
    "            sample_x.append(tx)\n",
    "            sample_y.append(ty)\n",
    "    # 别忘了起点和目标点\n",
    "    sample_x.append(sx)\n",
    "    sample_y.append(sy)\n",
    "    sample_x.append(gx)\n",
    "    sample_y.append(gy)\n",
    "\n",
    "    return sample_x, sample_y\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 主函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " start!!\n",
      "goal is found!\n"
     ]
    }
   ],
   "source": [
    "def main(rng=None):\n",
    "    print(\" start!!\")\n",
    "    fig = plt.figure(1)\n",
    "\n",
    "    # camara = Camera(fig)  # 保存动图时使用\n",
    "    camara = None\n",
    "    # start and goal position\n",
    "    sx = 10.0  # [m]\n",
    "    sy = 10.0  # [m]\n",
    "    gx = 50.0  # [m]\n",
    "    gy = 50.0  # [m]\n",
    "    robot_size = 5.0  # [m]\n",
    "\n",
    "    ox = []\n",
    "    oy = []\n",
    "\n",
    "    for i in range(60):\n",
    "        ox.append(i)\n",
    "        oy.append(0.0)\n",
    "    for i in range(60):\n",
    "        ox.append(60.0)\n",
    "        oy.append(i)\n",
    "    for i in range(61):\n",
    "        ox.append(i)\n",
    "        oy.append(60.0)\n",
    "    for i in range(61):\n",
    "        ox.append(0.0)\n",
    "        oy.append(i)\n",
    "    for i in range(40):\n",
    "        ox.append(20.0)\n",
    "        oy.append(i)\n",
    "    for i in range(40):\n",
    "        ox.append(40.0)\n",
    "        oy.append(60.0 - i)\n",
    "\n",
    "    if show_animation:\n",
    "        plt.plot(ox, oy, \".k\")\n",
    "        plt.plot(sx, sy, \"^r\")\n",
    "        plt.plot(gx, gy, \"^c\")\n",
    "        plt.grid(True)\n",
    "        plt.axis(\"equal\")\n",
    "        if camara != None:\n",
    "            camara.snap()\n",
    "\n",
    "    rx, ry = prm_planning(sx, sy, gx, gy, ox, oy,\n",
    "                          robot_size, camara=camara, rng=rng)\n",
    "\n",
    "    assert rx, 'Cannot found path'\n",
    "\n",
    "    if show_animation:\n",
    "        plt.plot(rx, ry, \"-r\")\n",
    "        plt.pause(0.001)\n",
    "        if camara != None:\n",
    "            camara.snap()\n",
    "            animation = camara.animate()\n",
    "            animation.save('trajectory.gif')\n",
    "        plt.savefig(\"result.png\")\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "\n",
    "main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.12 ('gobigger')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "0c7484b3574347463e16b31029466871583b0d4e5c4ad861e8848f2d3746b4de"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
